import os

from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain.output_parsers import ResponseSchema, StructuredOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda


# 调试输出
def debug_print(text):
    print("中间输出:", text)
    return text


if __name__ == '__main__':
    load_dotenv(override=True)
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")

    debug_node = RunnableLambda(debug_print)
    news_gen_prompt = PromptTemplate.from_template(
        "请根据以下新闻标题撰写一段简短的新闻内容(100字以内):\n\n标题: {title}"
    )
    # 第一个子链生成新闻内容
    news_chain = news_gen_prompt | model
    # 第二步 从正文提取结构化字段
    schemas = [
        ResponseSchema(name="time", description="事件发生的时间"),
        ResponseSchema(name="location", description="事件发生的地点"),
        ResponseSchema(name="event", description="发生的具体内容")
    ]
    parser = StructuredOutputParser.from_response_schemas(schemas)

    summary_prompt = PromptTemplate.from_template(
        "请根据以下新闻内容生成一个摘要:\n\n{news2}\n\n{format_instructions}"
    )
    # 第二个子链生成摘要
    summary_chain = (summary_prompt.partial(format_instructions=parser.get_format_instructions()) | model | parser)
    # 将两个子链组合成一个完整的链
    full_chain = news_chain | debug_node | summary_chain
    result = full_chain.invoke({"title": "中国和日本开始进行贸易合作"})
    print(result)
